using Newtonsoft.Json.Linq;
using SPI.AgentServer.Services.Execution;

namespace SPI.AgentServer.Services.Tools;

public class RunPrebuiltQueryTool : IQueryTool
{
    private readonly IDatabaseExecutor _executor;
    private readonly ILogger<RunPrebuiltQueryTool> _logger;
    private readonly Dictionary<string, PrebuiltQuery> _queries;

    public string Name => "run_prebuilt_query";
    public string Description => "Run predefined complex queries (for partitioned tables, cross-database queries, etc.)";

    public RunPrebuiltQueryTool(IDatabaseExecutor executor, ILogger<RunPrebuiltQueryTool> logger)
    {
        _executor = executor;
        _logger = logger;
        _queries = InitializePrebuiltQueries();
    }

    public async Task<ToolResult> ExecuteAsync(Dictionary<string, object> parameters, CancellationToken cancellationToken = default)
    {
        if (!ValidateParameters(parameters, out var errors))
        {
            return new ToolResult
            {
                Success = false,
                Errors = errors
            };
        }

        var queryName = parameters["queryName"].ToString()!;
        var queryParams = parameters.ContainsKey("parameters") ? 
            JObject.FromObject(parameters["parameters"]).ToObject<Dictionary<string, object>>() : 
            new Dictionary<string, object>();

        try
        {
            if (!_queries.TryGetValue(queryName, out var prebuiltQuery))
            {
                return new ToolResult
                {
                    Success = false,
                    Errors = new List<string> { $"Query '{queryName}' not found" }
                };
            }

            // Validate query parameters
            var paramErrors = ValidateQueryParameters(prebuiltQuery, queryParams ?? new Dictionary<string, object>());
            if (paramErrors.Any())
            {
                return new ToolResult
                {
                    Success = false,
                    Errors = paramErrors
                };
            }

            // Build the SQL with parameters
            var sql = BuildParameterizedSql(prebuiltQuery.SqlTemplate, queryParams);

            // Execute across multiple databases if needed
            if (prebuiltQuery.RequiresMultiDb)
            {
                var results = new List<object>();
                foreach (var db in prebuiltQuery.Databases)
                {
                    var result = await _executor.ExecuteQueryAsync(sql, db, cancellationToken);
                    results.Add(new { database = db, data = result.Data });
                }

                return new ToolResult
                {
                    Success = true,
                    Data = results,
                    Metadata = new Dictionary<string, object>
                    {
                        ["queryName"] = queryName,
                        ["databases"] = prebuiltQuery.Databases
                    }
                };
            }
            else
            {
                var result = await _executor.ExecuteQueryAsync(sql, prebuiltQuery.Databases.FirstOrDefault(), cancellationToken);
                return new ToolResult
                {
                    Success = true,
                    Data = result.Data,
                    Metadata = new Dictionary<string, object>
                    {
                        ["queryName"] = queryName,
                        ["rowCount"] = result.RowCount,
                        ["executionTimeMs"] = result.ExecutionTimeMs
                    }
                };
            }
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error executing prebuilt query: {QueryName}", queryName);
            return new ToolResult
            {
                Success = false,
                Errors = new List<string> { $"Execution error: {ex.Message}" }
            };
        }
    }

    public bool ValidateParameters(Dictionary<string, object> parameters, out List<string> errors)
    {
        errors = new List<string>();

        if (!parameters.ContainsKey("queryName") || string.IsNullOrWhiteSpace(parameters["queryName"]?.ToString()))
        {
            errors.Add("Parameter 'queryName' is required");
        }

        return errors.Count == 0;
    }

    private Dictionary<string, PrebuiltQuery> InitializePrebuiltQueries()
    {
        return new Dictionary<string, PrebuiltQuery>
        {
            ["user_activity_summary"] = new PrebuiltQuery
            {
                Name = "user_activity_summary",
                Description = "Get user activity summary across time periods",
                SqlTemplate = @"
                    SELECT 
                        DATE(activity_time) as activity_date,
                        COUNT(DISTINCT user_id) as unique_users,
                        COUNT(*) as total_activities
                    FROM user_activities_{partition}
                    WHERE activity_time BETWEEN @startDate AND @endDate
                    GROUP BY DATE(activity_time)
                    ORDER BY activity_date",
                Parameters = new Dictionary<string, string>
                {
                    ["startDate"] = "Start date (YYYY-MM-DD)",
                    ["endDate"] = "End date (YYYY-MM-DD)"
                },
                RequiresPartitioning = true
            },
            ["cross_db_sales_report"] = new PrebuiltQuery
            {
                Name = "cross_db_sales_report",
                Description = "Sales report across multiple regional databases",
                SqlTemplate = @"
                    SELECT 
                        product_id,
                        SUM(quantity) as total_quantity,
                        SUM(amount) as total_amount
                    FROM sales
                    WHERE sale_date BETWEEN @startDate AND @endDate
                    GROUP BY product_id",
                Parameters = new Dictionary<string, string>
                {
                    ["startDate"] = "Start date",
                    ["endDate"] = "End date"
                },
                RequiresMultiDb = true,
                Databases = new List<string> { "Sales_North", "Sales_South", "Sales_East", "Sales_West" }
            }
        };
    }

    private List<string> ValidateQueryParameters(PrebuiltQuery query, Dictionary<string, object> providedParams)
    {
        var errors = new List<string>();
        
        foreach (var requiredParam in query.Parameters.Keys)
        {
            if (!providedParams.ContainsKey(requiredParam))
            {
                errors.Add($"Required parameter '{requiredParam}' is missing");
            }
        }

        return errors;
    }

    private string BuildParameterizedSql(string template, Dictionary<string, object> parameters)
    {
        var sql = template;
        
        foreach (var param in parameters)
        {
            // Handle partition placeholders
            if (param.Key == "partition" && param.Value is DateTime date)
            {
                var partition = date.ToString("yyyyMM");
                sql = sql.Replace("{partition}", partition);
            }
            else
            {
                // Replace parameter placeholders
                sql = sql.Replace($"@{param.Key}", $"'{param.Value}'");
            }
        }

        return sql;
    }

    private class PrebuiltQuery
    {
        public string Name { get; set; } = string.Empty;
        public string Description { get; set; } = string.Empty;
        public string SqlTemplate { get; set; } = string.Empty;
        public Dictionary<string, string> Parameters { get; set; } = new();
        public bool RequiresPartitioning { get; set; }
        public bool RequiresMultiDb { get; set; }
        public List<string> Databases { get; set; } = new();
    }
}
